import torch
import torch.nn as nn


# class InConv(nn.Module):
#     def __init__(self, in_ch, out_ch): # 两次卷积
#         super(InConv, self).__init__()
#         self.conv = DoubleConv(in_ch, out_ch)
#
#     def forward(self, x):
#         x = self.conv(x)
#         return x
# InConv 类继承自 PyTorch 的 nn.Module 类
class InConv(nn.Module):
    # 类的初始化函数
    # in_ch 表示输入通道数，out_ch 表示输出通道数
    def __init__(self, in_ch, out_ch):
        # 调用 nn.Module 的构造函数进行初始化
        super(InConv, self).__init__()
        # 创建一个 DoubleConv 模块实例，DoubleConv 是一个假定已定义的模块
        # 通常，这个模块会进行两次卷积操作，可能还包括激活函数和批量归一化
        self.conv = DoubleConv(in_ch, out_ch)

    # 定义前向传播函数，接收输入特征图 x
    def forward(self, x):
        # 将输入 x 传递给之前定义的双卷积模块
        x = self.conv(x)
        # 返回经过两次卷积操作后的特征图
        return x

class Down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(Down, self).__init__()
        self.mpconv = nn.Sequential( # 可以作为一个复合层一起使用，先缩小特征图的空间尺寸，然后进行一次或多次卷积来提取特征
            nn.MaxPool3d(2, 2),
            DoubleConv(in_ch, out_ch)
        )

    def forward(self, x):
        x = self.mpconv(x)
        return x

# class OutConv(nn.Module):
#     def __init__(self, in_ch, out_ch):
#         super(OutConv, self).__init__()
#         self.conv = nn.Conv3d(in_ch, out_ch, 1)
#         # self.sigmoid = nn.Sigmoid()
#
#     def forward(self, x):
#         x = self.conv(x)
#         # x = self.sigmoid(x)
#         return x
class OutConv(nn.Module):
    # 类的初始化函数，接收输入通道数 in_ch 和输出通道数 out_ch
    def __init__(self, in_ch, out_ch):
        # 调用 nn.Module 的构造函数，进行必要的初始化
        super(OutConv, self).__init__()
        # 创建一个卷积层，用于将特征图的通道数从 in_ch 改变到 out_ch
        # 这里使用的是 3D 卷积，卷积核大小设为 1，这通常用于改变通道数而不增加空间维度
        self.conv = nn.Conv3d(in_ch, out_ch, 1)

        # 下一行被注释掉了，如果去掉注释，会创建一个 Sigmoid 激活函数层
        # 激活函数通常用于分类任务的输出层，将值压缩到 [0, 1] 范围内
        # self.sigmoid = nn.Sigmoid()

    # 定义前向传播函数，接收输入特征图 x
    def forward(self, x):
        # 将输入 x 传递给之前定义的卷积层
        x = self.conv(x)

        # 下一行被注释掉了，如果去掉注释，会将卷积后的结果通过 Sigmoid 函数进行激活
        # x = self.sigmoid(x)

        # 返回经过卷积（和可选的 Sigmoid 激活）后的特征图
        return x

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        # nn.Conv3d 层的作用是学习输入数据的三维特征。它常用于处理三维数据，例如体积医学图像或视频序列，它可以捕获这些数据中的空间和时间特征。在U-Net架构或其他3D CNN架构中，这些层是构建深度学习模型的基本构件
        self.conv = nn.Sequential(
            nn.Conv3d(in_ch, out_ch, kernel_size=3, stride=1, padding=1),    # 3d 卷积层 步长为1
            nn.BatchNorm3d(out_ch),  # 第一个批归一化层
            nn.ReLU(inplace=True),  # 激活函数
            nn.Conv3d(out_ch, out_ch, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm3d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x

class Up(nn.Module):
    # def __init__(self, in_ch, skip_ch,out_ch):
    #     super(Up, self).__init__()
    #     self.up = nn.ConvTranspose3d(in_ch, in_ch, kernel_size=2, stride=2) # 转置卷积 插值法也可
    #     self.conv = DoubleConv(in_ch+skip_ch, out_ch)

    # 类的初始化函数
    # in_ch 表示输入通道数，skip_ch 表示从跳跃连接来的旁路输入的通道数，out_ch 表示输出通道数
    def __init__(self, in_ch, skip_ch, out_ch):
        # 调用 nn.Module 的构造函数进行初始化
        super(Up, self).__init__()
        # 创建一个转置卷积层（逆卷积层），用于上采样（增加特征图的尺寸）
        # in_ch 是转置卷积层的输入通道数，也是输出通道数（因为这里是先上采样再进行拼接）
        # kernel_size 是卷积核的大小，stride 是步长，这些参数共同决定了特征图的上采样倍数
        self.up = nn.ConvTranspose3d(in_ch, in_ch, kernel_size=2, stride=2)

        # 创建一个 DoubleConv 模块实例，用于执行两次卷积操作
        # 输入通道数是 in_ch（从上采样得到的）加上 skip_ch（从跳跃连接得到的），输出通道数是 out_ch
        self.conv = DoubleConv(in_ch + skip_ch, out_ch)

    def forward(self, x1, x2):
        # 通过转置卷积层对 x1 进行上采样
        x1 = self.up(x1)
        # 将上采样后的结果 x1 与特征图 x2 进行通道维度上的拼接
        # dim=1 表示在通道维度上进行拼接
        x = torch.cat([x2, x1], dim=1)
        # 将拼接后的特征图传递给双重卷积层进行进一步的卷积操作
        x = self.conv(x)
        # 返回卷积后的特征图
        return x

#
# class UNet(nn.Module):
#     def __init__(self, in_channels, num_classes):
#         super(UNet, self).__init__()
#         features = [32,64,128,256]
#
#         self.inc = InConv(in_channels, features[0])
#         self.down1 = Down(features[0], features[1])
#         self.down2 = Down(features[1], features[2])
#         self.down3 = Down(features[2], features[3])
#         self.down4 = Down(features[3], features[3])
#
#         self.up1 = Up(features[3], features[3], features[2])
#         self.up2 = Up(features[2], features[2], features[1])
#         self.up3 = Up(features[1], features[1], features[0])
#         self.up4 = Up(features[0], features[0], features[0])
#         self.outc = OutConv(features[0], num_classes)
#
#     def forward(self, x):
#         x1 = self.inc(x)
#         x2 = self.down1(x1)
#         x3 = self.down2(x2)
#         x4 = self.down3(x3)
#         x5 = self.down4(x4)
#
#         x = self.up1(x5, x4)
#         x = self.up2(x, x3)
#         x = self.up3(x, x2)
#         x = self.up4(x, x1)
#         x = self.outc(x)
#         return x


class UNet(nn.Module):
    def __init__(self, in_channels, num_classes):
        # 调用 nn.Module 的构造函数进行初始化
        super(UNet, self).__init__()

        # 定义网络的特征层数组，这里使用了 4 个不同的特征层级
        features = [32, 64, 128, 256]

        # 初始化 U-Net 网络的各个模块
        self.inc = InConv(in_channels, features[0])  # 输入卷积块
        self.down1 = Down(features[0], features[1])  # 第一个下采样块
        self.down2 = Down(features[1], features[2])  # 第二个下采样块
        self.down3 = Down(features[2], features[3])  # 第三个下采样块
        self.down4 = Down(features[3], features[3])  # 第四个下采样块

        # 上采样部分，逐步增加特征图的空间维度，并减少特征通道数
        self.up1 = Up(features[3], features[3], features[2])  # 第一个上采样块
        self.up2 = Up(features[2], features[2], features[1])  # 第二个上采样块
        self.up3 = Up(features[1], features[1], features[0])  # 第三个上采样块
        self.up4 = Up(features[0], features[0], features[0])  # 第四个上采样块

        # 输出卷积块，将最后的特征图转换为类分数图
        self.outc = OutConv(features[0], num_classes)

    def forward(self, x):
        # 依次通过初始化的各个模块
        x1 = self.inc(x)  # 输入数据通过输入卷积块
        x2 = self.down1(x1)  # 下采样
        x3 = self.down2(x2)  # 下采样
        x4 = self.down3(x3)  # 下采样
        x5 = self.down4(x4)  # 下采样

        # 上采样，并通过跳跃连接与对应的下采样特征图进行拼接
        x = self.up1(x5, x4)  # 上采样并与 x4 拼接
        x = self.up2(x, x3)  # 上采样并与 x3 拼接
        x = self.up3(x, x2)  # 上采样并与 x2 拼接
        x = self.up4(x, x1)  # 上采样并与 x1 拼接

        # 最后通过输出卷积块得到最终的类分数图
        x = self.outc(x)
        return x  # 返回最终的输出

if __name__ == '__main__':
    # 这行代码使用PyTorch库创建一个形状为(1, 4, 160, 160, 128)
    # 的随机张量x。张量的维度解释如下：
    #
    # 1
    # 是批次大小（batch
    # size），表示一次通过网络的数据样本数。
    # 4
    # 是输入通道数，表示每个数据样本有4个通道，这对应于in_channels = 4。
    # 160, 160, 128
    # 是数据样本的空间维度，这里看起来是一个3D数据的格式（可能是一个医学影像中的体积数据）。
    x = torch.randn(1, 4, 160, 160, 128)
    net = UNet(in_channels=4, num_classes=4)
    y = net(x)
    print("params: ", sum(p.numel() for p in net.parameters()))
    print(y.shape)